"""
基于 gRPC 的机械臂 Gymnasium 环境封装

将 StreamClient (gRPC客户端) 封装为符合 Gymnasium 接口的环境类。
"""

from __future__ import annotations

from typing import Any, Dict, Optional, Tuple

import gymnasium as gym
import numpy as np
from gymnasium import spaces

from parnassus_train.clients.arm_stream.client import StreamClient


class GrpcArmEnv(gym.Env):
    """
    基于 gRPC 的机械臂 Gymnasium 环境
    
    观测空间：4维连续空间 [q0, q1, dq0, dq1]
    - q0: 关节0的角度
    - q1: 关节1的角度（摆杆相对角度）
    - dq0: 关节0的角速度
    - dq1: 关节1的角速度
    
    动作空间：1维连续空间 [-1, 1]
    - 关节0的力矩控制
    
    通过 gRPC 与远程环境服务器通信。
    """
    
    metadata = {"render_modes": ["human"], "render_fps": 30}
    
    def __init__(
        self,
        server_address: str = "localhost:50051",
        obs_dim: int = 4,
        action_dim: int = 1,
        max_episode_steps: int = 500,
        reward_shaping: bool = True,
        render_mode: Optional[str] = None,
    ):
        """
        初始化环境
        
        Args:
            server_address: gRPC 服务器地址
            obs_dim: 观测空间维度
            action_dim: 动作空间维度
            max_episode_steps: 每个episode的最大步数
            reward_shaping: 是否使用奖励塑形
            render_mode: 渲染模式（暂不支持）
        """
        super().__init__()
        
        self.server_address = server_address
        self.max_episode_steps = max_episode_steps
        self.reward_shaping = reward_shaping
        self.render_mode = render_mode
        
        # 定义观测空间和动作空间
        self.observation_space = spaces.Box(
            low=-np.inf,
            high=np.inf,
            shape=(obs_dim,),
            dtype=np.float32,
        )
        
        self.action_space = spaces.Box(
            low=-1.0,
            high=1.0,
            shape=(action_dim,),
            dtype=np.float32,
        )
        
        # gRPC 客户端（延迟初始化）
        self._client: Optional[StreamClient] = None
        self._current_step = 0
        self._episode_seed = None
        
    def _ensure_client(self):
        """确保 gRPC 客户端已连接"""
        if self._client is None:
            self._client = StreamClient(self.server_address)
            self._client.__enter__()
    
    def _compute_shaped_reward(
        self,
        observation: np.ndarray,
        env_reward: float,
        terminated: bool,
    ) -> float:
        """
        计算塑形后的奖励
        
        目标：让摆杆（关节1）竖直向上并保持稳定
        
        Args:
            observation: 当前观测 [q0, q1, dq0, dq1]
            env_reward: 环境原始奖励
            terminated: 是否终止
            
        Returns:
            shaped_reward: 塑形后的奖励
        """
        if not self.reward_shaping:
            return env_reward
        
        q0, q1, dq0, dq1 = observation
        
        # 1. 竖直向上的角度奖励（核心目标）
        # 将角度归一化到[-π, π]范围
        q1_normalized = np.arctan2(np.sin(q1), np.cos(q1))
        angle_to_upright = abs(q1_normalized)  # 距离竖直向上的角度差
        angle_reward = -angle_to_upright  # 角度越接近0，奖励越高
        
        # 2. 低速度奖励（稳定性）
        velocity_penalty = -0.1 * (dq0**2 + dq1**2)  # 速度越小越好
        
        # 3. 存活奖励（鼓励持续平衡）
        survival_reward = 1.0 if not terminated else 0.0
        
        # 4. 额外的竖直位置奖励（当角度很小时给予大奖励）
        if angle_to_upright < 0.2:  # 约11.5度以内
            upright_bonus = 5.0 * (0.2 - angle_to_upright)
        else:
            upright_bonus = 0.0
        
        # 组合奖励
        reward = (
            angle_reward * 10.0 +      # 角度是最重要的
            velocity_penalty +          # 抑制过大的速度
            survival_reward +           # 基础存活奖励
            upright_bonus               # 竖直位置额外奖励
        )
        
        return reward
    
    def reset(
        self,
        seed: Optional[int] = None,
        options: Optional[Dict[str, Any]] = None,
    ) -> Tuple[np.ndarray, Dict[str, Any]]:
        """
        重置环境
        
        Args:
            seed: 随机种子
            options: 额外选项
            
        Returns:
            observation: 初始观测
            info: 附加信息
        """
        super().reset(seed=seed)
        
        # 确保客户端已连接
        self._ensure_client()
        
        # 重置环境（通过 gRPC）
        self._episode_seed = seed if seed is not None else np.random.randint(0, 100000)
        observation = self._client.reset(seed=self._episode_seed)
        
        # 转换为 numpy 数组
        observation = np.array(observation, dtype=np.float32)
        
        # 重置步数计数器
        self._current_step = 0
        
        info = {
            "episode_seed": self._episode_seed,
        }
        
        return observation, info
    
    def step(
        self,
        action: np.ndarray,
    ) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
        """
        执行一步动作
        
        Args:
            action: 动作数组 [action_dim]
            
        Returns:
            observation: 新观测
            reward: 奖励
            terminated: 是否因任务完成/失败而终止
            truncated: 是否因超时等原因截断
            info: 附加信息
        """
        # 确保客户端已连接
        self._ensure_client()
        
        # 确保 action 是正确的类型
        if not isinstance(action, np.ndarray):
            action = np.array(action, dtype=np.float32)
        
        # 执行动作（通过 gRPC）
        step_reply = self._client.step(action)
        
        # 解析返回结果
        observation = np.array(step_reply["observation"], dtype=np.float32)
        env_reward = step_reply["reward"]
        terminated = step_reply["terminated"]
        truncated = step_reply["truncated"]
        
        # 处理可能的 None 值
        if env_reward is None:
            env_reward = 0.0
        if terminated is None:
            terminated = False
        if truncated is None:
            truncated = False
        
        # 计算塑形后的奖励
        reward = self._compute_shaped_reward(observation, env_reward, terminated)
        
        # 更新步数
        self._current_step += 1
        
        # 检查是否达到最大步数
        if self._current_step >= self.max_episode_steps:
            truncated = True
        
        # 附加信息
        info = {
            "env_reward": env_reward,  # 原始环境奖励
            "shaped_reward": reward,    # 塑形后的奖励
            "step": self._current_step,
        }
        
        return observation, reward, terminated, truncated, info
    
    def render(self):
        """渲染环境（暂不支持）"""
        if self.render_mode == "human":
            # 可以在这里实现可视化
            pass
    
    def close(self):
        """关闭环境并断开 gRPC 连接"""
        if self._client is not None:
            self._client.__exit__(None, None, None)
            self._client = None
    
    def __enter__(self):
        """上下文管理器入口"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.close()
